package njust.csc.aecore.judge

import io.github.oshai.kotlinlogging.KotlinLogging
import kotlinx.serialization.encodeToString
import kotlinx.serialization.json.Json
import njust.csc.aecore.entity.*
import njust.csc.aecore.utils.SerializeUtils
import java.io.BufferedReader
import java.io.DataOutputStream
import java.io.InputStreamReader
import java.net.HttpURLConnection
import java.net.URL
import java.nio.charset.StandardCharsets
private val logger = KotlinLogging.logger {}

val prompt = """
# 角色
你是一个公安机关工作人员，是手机app反诈专家，能够从给定的信息中判断一个apk是否涉嫌诈骗。

## 技能
- 涉诈 APK 通联地址 URL 可能是无规律的，不具备明确的相关性或意义，也可能涉及金钱等相关敏感词汇
- 涉诈 APK 的名称可能模仿知名应用，使用相似或误导性的文字拼写，比如在品牌名称中加入特殊字符、数字或轻微变形
- 涉诈 APK 的包名可能包含乱码或随机字符，或模仿合法应用的包名但有细微差异
- 签名上，可能使用无效或自签名，签名名称是乱码，或者盗用已知的被撤销的签名，缺乏信誉度和安全性验证。
- 版本号可能随意设置，与实际内容不符，有时故意设置为高版本以吸引用户更新或显得更“新”。
- 涉诈应用可能会要求过于多的权限，尤其是敏感权限

## 限制
- 你会得到一个 json 字符串，包含一个 apk 的所有信息，其中certification为签名信息，urls是apk的可能会请求的服务端地址，permissions为权限信息，packageName为包名，versionInfo为版本信息
- 不要倾向于揣测 APK 涉诈
- 请根据给定的 json 信息，在回答的开头给出一个范围为 0 到 100 的数字，表示该 APK 涉诈的概率。不要输出其他任何文本内容
```
请分析如下apk是否涉诈
%s
```
""".trim()

/**
 * @Author ZhouHang
 * @Date 24/6/23 周日
 */
class QWenJudge : LLMJudge {
    private val urlStr = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation"
    private val apiKey = System.getenv("QWEN_API_KEY")
    private val url = URL(urlStr)

    override fun ask(apkInfo: String): String {
        val apk = SerializeUtils.load(apkInfo)
        return ask(apk)
    }

    override fun ask(apk: APK): String {
        // 删除 apk 中的一些字段, 专门给大模型使用
        val apkSimple = APKSimple(
            name = apk.name,
            packageName = apk.packageName,
            size = apk.size,
            versionInfo = apk.versionInfo,
            permissions = apk.permissions,
            subject = apk.v1SignatureInfo.certification.subject,
            signatureType = apk.v1SignatureInfo.certification.signatureType,
            md5 = apk.v1SignatureInfo.certification.md5,
            urls = apk.urls
        )

        val apkSimpleInfo = Json.encodeToString(apkSimple)

        val connection = url.openConnection() as HttpURLConnection
        connection.setRequestMethod("POST")
        connection.setRequestProperty("Content-Type", "application/json")
        connection.setRequestProperty("Authorization", "Bearer $apiKey")
        connection.setDoOutput(true)

        val body = QWenRequestBody(
            model = "qwen-max",
            input = QWenInput(
                messages = listOf(
                    Message(
                        role = Role.SYSTEM,
                        content = "You are a helpful assistant"
                    ),
                    Message(
                        role = Role.USER,
                        content = prompt.format(apkSimpleInfo)
                    )
                )
            ),
            parameters = QWenParameter(
                result_format = "message"
            )
        )

        logger.info { Json.encodeToString(body) }
        DataOutputStream(connection.outputStream).use { wr ->
            wr.write(Json.encodeToString(body).toByteArray(StandardCharsets.UTF_8))
            wr.flush()
        }
        val response = StringBuilder()
        BufferedReader(
            InputStreamReader(connection.inputStream)
        ).use { `in` ->
            var inputLine: String?
            while ((`in`.readLine().also { inputLine = it }) != null) {
                response.append(inputLine)
            }
        }
        logger.info { response }
        val resp = Json.decodeFromString<QWenResponse>(response.toString())
        connection.disconnect()
        return resp.output.choices[0].message.content
    }
}
